package com.cetc.sdp.kmga.cs.stream;

import com.cetc.sdp.kmga.cs.audit.AuditBriefInfo;
import com.cetc.sdp.kmga.cs.audit.OpCode;
import com.cetc.sdp.kmga.cs.jdbc.AuditDAO;
import com.cetc.sdp.kmga.cs.util.Tool;
import com.rabbitmq.client.QueueingConsumer;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.api.java.JavaInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.rabbitmq.RabbitMQUtils;
import org.apache.spark.streaming.rabbitmq.distributed.JavaRabbitMQDistributedKey;
import org.apache.spark.streaming.rabbitmq.models.ExchangeAndRouting;

import java.time.format.DateTimeFormatter;
import java.time.format.FormatStyle;
import java.util.*;

/**
 * 采用RabbitMQ 作为Spark Streaming 输入源
 * @description:
 * @author： DengQiang
 * @date: 2017/10/19 10:12
 */
public class RabbitMqStream {

    private static final int ARGS_LEN = 3;

    private static final String CONF_FILE_PATH = "RabbitmqStream.properties";

    public static void main(String[] args) throws InterruptedException {
        if (args.length < ARGS_LEN) {
            System.out.println("Usage：<Application Name> <Batch duration> <Queue>");
            System.exit(1);
        }
        String appName = args[0];
        int duration = Integer.parseInt(args[1]);
        String queue = args[2];

        Properties prop = Tool.loadProperties(CONF_FILE_PATH);
        Map<String, String> rabbitmqParams = new HashMap<>(32);
        prop.entrySet().forEach(entry -> rabbitmqParams.put(entry.getKey().toString(), entry.getValue().toString()));

        //初始化上下文环境
        SparkConf conf = new SparkConf().setAppName(appName);
        JavaSparkContext jsc = new JavaSparkContext(conf);
        jsc.setLogLevel("ERROR");
        JavaStreamingContext jssc = new JavaStreamingContext(jsc, Durations.seconds(duration));

        List<JavaRabbitMQDistributedKey> distributedKeys = new LinkedList<>();
        distributedKeys.add(new JavaRabbitMQDistributedKey(queue,
                new ExchangeAndRouting(""),
                rabbitmqParams)
        );
        Function<QueueingConsumer.Delivery, String> messageHandler = (Function<QueueingConsumer.Delivery, String>) message -> new String(message.getBody());
        JavaInputDStream<String> rabbitmqStream =
                RabbitMQUtils.createJavaDistributedStream(jssc, String.class, distributedKeys, rabbitmqParams, messageHandler);

        rabbitmqStream.foreachRDD(stringJavaRDD -> {
            stringJavaRDD.foreachPartition(stringIterator -> {
                List<AuditBriefInfo> res = new ArrayList<>(64);
                while (stringIterator.hasNext()) {
                    String event = stringIterator.next();
                    String data = event.substring(17, event.length()-2);
                    StringTokenizer tokenizer = new StringTokenizer(data, ";");
                    Map<String, String> dmap = new HashMap<>(10);
                    while (tokenizer.hasMoreTokens()) {
                        String token = tokenizer.nextToken();
                        String[] entry = token.split("=");
                        if (entry.length > 2) {
                            dmap.put(entry[0], entry[1]);
                        }
                    }

                    AuditBriefInfo briefInfo = new AuditBriefInfo();
                    /**
                     * 判断流量方向
                     */
                    OpCode opCode = getOpCode(dmap.get("accessType"));
                    if (opCode != null) {
                        briefInfo.setOperTypeCode(opCode.getOperTypeCode());
                    }
                    DateTimeFormatter.ofPattern("E M d HH:mm:ss z yy").parse("Thu Oct 19 16:00:07 EDT 2017");
                }

            });
        });

//        rabbitmqStream.print();


        jssc.start();
        jssc.awaitTermination();
    }

    private volatile static Map<String, OpCode> opCodeMap;

    public static OpCode getOpCode(String oper) {
        if (opCodeMap ==null) {
            synchronized (RabbitMqStream.class) {
                if (opCodeMap == null) {
                    opCodeMap = AuditDAO.findAllOpCode();
                }
            }
        }
        return opCodeMap.get(oper);
    }
}
